from crazyflow.sim import Sim
from crazyflow.sim.data import SimData

sim = Sim()

print(sim.step_pipeline)                # Pipeline(step_attitude_controller -> ... -> clip_floor_pos)
sim.step_pipeline.names                 # ordered stage names
"integration" in sim.step_pipeline      # True
sim.step_pipeline.index("integration")  # position of a stage
sim.step_pipeline["clip_floor_pos"]     # look up a stage's function by name
len(sim.step_pipeline)                  # number of stages

def disturbance(data: SimData) -> SimData:
    return data.replace(states=data.states.replace(vel=data.states.vel + 1e-5))

def monitor(data: SimData) -> SimData:
    return data

sim.step_pipeline.insert_before("integration", disturbance)
sim.step_pipeline.insert_after("integration", monitor)
sim.step_pipeline.prepend(monitor, name="pre_monitor")          # name= for a reused callable
sim.step_pipeline.replace("clip_floor_pos", lambda data: data)  # swap impl, keep name + position
sim.step_pipeline.remove("monitor")

sim.step_pipeline = sim.step_pipeline + monitor  # Same as append with __name__
sim.step_pipeline.extend(some_other_pipeline)    # Adding multiple steps at once
for stage in sim.step_pipeline:                  # iterating yields the functions, in order
    ...